import onnxruntime as ort
import numpy as np


# 定义加载和预处理数据的函数
def preprocess_inputs(des, tweet, num_prop, cat_prop, edge_index, edge_type):
    """
    根据模型输入格式预处理输入数据。
    """
    inputs = {
        "des": des.astype(np.float32),
        "tweet": tweet.astype(np.float32),
        "num_prop": num_prop.astype(np.float32),
        "cat_prop": cat_prop.astype(np.float32),
        "edge_index": edge_index.astype(np.int64),
        "edge_type": edge_type.astype(np.int64),
    }
    return inputs


# 定义推理函数
def infer_onnx_model(model_path, des, tweet, num_prop, cat_prop, edge_index, edge_type):
    """
    加载ONNX模型并进行推理。
    """
    # 加载ONNX模型
    session = ort.InferenceSession(model_path)

    # 获取模型输入名
    input_names = [input.name for input in session.get_inputs()]
    output_names = [output.name for output in session.get_outputs()]
    print("Input names:", input_names)
    print("Output names:", output_names)

    # 预处理输入
    inputs = preprocess_inputs(des, tweet, num_prop, cat_prop, edge_index, edge_type)

    # 推理
    results = session.run(output_names, inputs)
    return results


# 示例数据（替换为你的真实数据）
des = np.random.randn(1, 768)  # 替换为实际的用户描述特征
tweet = np.random.randn(1, 768)  # 替换为用户推文特征
num_prop = np.random.randn(1, 5)  # 替换为用户的数值属性
cat_prop = np.random.randn(1, 3)  # 替换为用户的分类属性
edge_index = np.array([[0], [0]])  # 图中与该用户相关的边索引
edge_type = np.array([0])  # 边类型（可调整为用户相关的实际值）

# 模型路径
model_path = "model.onnx"

# 推理并输出结果
results = infer_onnx_model(model_path, des, tweet, num_prop, cat_prop, edge_index, edge_type)
print("Inference results:", results)
